from streamlit_cookies_controller import CookieController,RemoveEmptyElementContainer
from itsdangerous import URLSafeTimedSerializer
from time import sleep

class SecureCookieManager(CookieController):
    """安全的 Cookie 管理器"""
    
    def __init__(self, secret_key="hellokitty"):
        # 调用基类的初始化方法
        super().__init__()
        self.secret_key = secret_key
        self.serializer = URLSafeTimedSerializer(self.secret_key)
        RemoveEmptyElementContainer.remove_empty_element = True  # 开启移除空元素的功能

    def create_signed_cookie(self, key, value):
        """创建签名 Cookie"""
        if not key:
            return False # 空的 key 值不允许创建 Cookie
        else:
            signed_value = self.serializer.dumps(value)
            self.set(key, signed_value)
            print("Key:", key,"Signed Value:", signed_value)
            sleep(1)  # 等待 1 秒，确保 Cookie 被浏览器保存
            return True

    def verify_signed_cookie(self, signed_value):
        """获取签名 Cookie 的签名前的值"""
        return self.serializer.loads(signed_value) 

    def get_all_cookies(self):
        """获取所有 Cookie"""
        return self.getAll()  # 利用基类的 get_all 方法获取所有 Cookie

    def get_cookie(self, key):
        """获取指定的 Cookie"""
        signed_value = self.get(key)  # 利用基类的 get 方法获取指定的 Cookie
        if signed_value:
            return self.verify_signed_cookie(signed_value)  # 验证并返回 Cookie 的值
        else:
            return None

    def remove_cookie(self, key):
        """删除指定的 Cookie"""
        if self.get(key):  # 如果指定的 Cookie 存在，则删除之
            print("Removing Cookie:", key)
            self.remove(key)  # 利用基类的 remove 方法删除 Cookie
            sleep(1)  # 等待 1 秒，确保 Cookie 被浏览器删除
            return True
        else:
            print("Cookie not found:", key)
            return False  # 未找到指定的 Cookie，返回 False
        
# 示例用法
if __name__ == "__main__":
    pass
